# database/stock_database.py
import sqlite3
import pandas as pd
import os
import shutil
from datetime import datetime, timedelta
import logging
from typing import List, Dict, Optional
from sheng.stock.basic_data_v1.config.database_config import DatabaseConfig
from sqlalchemy.dialects.postgresql import Any


class StockDatabase:
    def __init__(self, config: DatabaseConfig = None):
        self.config = config or DatabaseConfig()
        self.db_path = self.config.get_db_path()
        self.logger = self._setup_logging()
        self._init_database()

    def _setup_logging(self):
        """设置日志"""
        logging.basicConfig(
            level=logging.INFO,
            format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
        )
        return logging.getLogger(__name__)

    def _init_database(self):
        """初始化数据库"""
        self.logger.info(f"初始化数据库: {self.db_path}")

        # 检查数据库文件是否存在
        db_exists = os.path.exists(self.db_path)

        conn = self._get_connection()
        try:
            cursor = conn.cursor()

            # 执行所有建表SQL
            table_scripts = self._get_table_scripts()
            for table_name, script in table_scripts.items():
                try:
                    # 分割SQL语句（处理多个语句）
                    statements = [s.strip() for s in script.split(';') if s.strip()]
                    for statement in statements:
                        if statement:  # 确保不是空字符串
                            cursor.execute(statement)
                    self.logger.info(f"表 {table_name} 初始化成功")
                except Exception as e:
                    self.logger.warning(f"表 {table_name} 可能已存在: {e}")

            conn.commit()

            if not db_exists:
                self.logger.info("新建数据库文件成功")
            else:
                self.logger.info("数据库连接成功")

        except Exception as e:
            self.logger.error(f"数据库初始化失败: {e}")
            raise
        finally:
            conn.close()

    def _get_table_scripts(self) -> Dict[str, str]:
        """返回完整的建表SQL语句"""
        return {
            "stock_basic": """
                CREATE TABLE IF NOT EXISTS stock_basic (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    symbol VARCHAR(20) NOT NULL UNIQUE,
                    name VARCHAR(100) NOT NULL,
                    market VARCHAR(10) NOT NULL,
                    industry VARCHAR(50),
                    listing_date DATE,
                    total_shares DECIMAL(18,2),
                    float_shares DECIMAL(18,2),
                    country VARCHAR(20) DEFAULT 'CN',
                    currency VARCHAR(10) DEFAULT 'CNY',
                    data_source VARCHAR(20) NOT NULL,
                    created_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    updated_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                )
            """,

            "daily_quotes": """
                CREATE TABLE IF NOT EXISTS daily_quotes (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    symbol VARCHAR(20) NOT NULL,
                    trade_date DATE NOT NULL,
                    open_price DECIMAL(12,4),
                    high_price DECIMAL(12,4),
                    low_price DECIMAL(12,4),
                    close_price DECIMAL(12,4),
                    volume DECIMAL(18,2),
                    turnover DECIMAL(18,2),
                    amplitude DECIMAL(8,4),
                    change_rate DECIMAL(8,4),
                    turnover_rate DECIMAL(8,4),
                    pe_ratio DECIMAL(10,4),
                    pb_ratio DECIMAL(10,4),
                    data_source VARCHAR(20) NOT NULL,
                    created_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    UNIQUE(symbol, trade_date, data_source)
                )
            """,

            "financial_data": """
                CREATE TABLE IF NOT EXISTS financial_data (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    symbol VARCHAR(20) NOT NULL,
                    report_date DATE NOT NULL,
                    report_type VARCHAR(10) NOT NULL,
                    eps DECIMAL(10,4),
                    revenue DECIMAL(18,2),
                    net_profit DECIMAL(18,2),
                    total_assets DECIMAL(18,2),
                    total_liabilities DECIMAL(18,2),
                    equity DECIMAL(18,2),
                    roe DECIMAL(8,4),
                    roa DECIMAL(8,4),
                    gross_margin DECIMAL(8,4),
                    net_margin DECIMAL(8,4),
                    data_source VARCHAR(20) NOT NULL,
                    created_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    UNIQUE(symbol, report_date, report_type, data_source)
                );
                CREATE INDEX IF NOT EXISTS idx_financial_symbol_date ON financial_data(symbol, report_date);
            """,

            # 分红数据表
            "dividend_data": """
                CREATE TABLE IF NOT EXISTS dividend_data (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    symbol VARCHAR(20) NOT NULL,
                    ex_dividend_date DATE NOT NULL,
                    dividend DECIMAL(10,4),
                    dividend_rate DECIMAL(8,4),
                    payout_ratio DECIMAL(8,4),
                    data_source VARCHAR(20) NOT NULL,
                    created_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    UNIQUE(symbol, ex_dividend_date, data_source)
                );
                CREATE INDEX IF NOT EXISTS idx_dividend_symbol_date ON dividend_data(symbol, ex_dividend_date);
            """,

            # 添加Baostock特有的财务指标表
            "baostock_indicators": """
                CREATE TABLE IF NOT EXISTS baostock_indicators (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    symbol VARCHAR(20) NOT NULL,
                    trade_date DATE NOT NULL,
                    pe_ratio DECIMAL(12,4),
                    pb_ratio DECIMAL(12,4),
                    ps_ratio DECIMAL(12,4),
                    pcf_ratio DECIMAL(12,4),
                    roe DECIMAL(8,4),
                    net_margin DECIMAL(8,4),
                    gross_margin DECIMAL(8,4),
                    equity_growth_rate DECIMAL(8,4),
                    revenue_growth_rate DECIMAL(8,4),
                    data_source VARCHAR(20) DEFAULT 'baostock',
                    created_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    UNIQUE(symbol, trade_date)
                );
                CREATE INDEX IF NOT EXISTS idx_bs_indicators_symbol_date ON baostock_indicators(symbol, trade_date);
            """,

            # 分红数据表（增强版）
            "dividend_data_enhanced": """
                CREATE TABLE IF NOT EXISTS dividend_data_enhanced (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    symbol VARCHAR(20) NOT NULL,
                    dividend_date DATE,
                    ex_dividend_date DATE,
                    dividend_per_share DECIMAL(10,4),
                    dividend_ratio DECIMAL(8,4),
                    update_date DATE,
                    data_source VARCHAR(20) NOT NULL,
                    created_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    UNIQUE(symbol, dividend_date, data_source)
                );
            """,

            # 数据更新日志表
            "update_log": """
                CREATE TABLE IF NOT EXISTS update_log (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    data_type VARCHAR(20) NOT NULL,
                    data_source VARCHAR(20) NOT NULL,
                    symbol VARCHAR(20),
                    start_date DATE,
                    end_date DATE,
                    records_updated INTEGER,
                    status VARCHAR(10) NOT NULL, -- SUCCESS, FAILED
                    error_message TEXT,
                    update_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                );
                CREATE INDEX IF NOT EXISTS idx_update_log_time ON update_log(update_time);
            """
        }

    def _get_connection(self) -> sqlite3.Connection:
        """获取数据库连接"""
        try:
            conn = sqlite3.connect(self.db_path)

            # 设置连接参数
            conn.execute("PRAGMA journal_mode=WAL")  # 写前日志模式
            conn.execute("PRAGMA synchronous=NORMAL")  # 平衡性能和数据安全
            conn.execute("PRAGMA cache_size=-64000")  # 设置缓存大小
            conn.execute("PRAGMA foreign_keys=ON")  # 启用外键约束

            return conn
        except Exception as e:
            self.logger.error(f"数据库连接失败: {e}")
            raise

    def get_database_info(self) -> Dict[str, Any]:
        """获取数据库信息"""
        conn = self._get_connection()
        try:
            cursor = conn.cursor()

            info = {}

            # 获取表信息
            cursor.execute("""
                SELECT name FROM sqlite_master 
                WHERE type='table' AND name NOT LIKE 'sqlite_%'
            """)
            tables = [row[0] for row in cursor.fetchall()]
            info['tables'] = tables

            # 获取各表记录数
            table_counts = {}
            for table in tables:
                cursor.execute(f"SELECT COUNT(*) FROM {table}")
                table_counts[table] = cursor.fetchone()[0]
            info['table_counts'] = table_counts

            # 获取数据库大小
            if os.path.exists(self.db_path):
                info['db_size_mb'] = round(os.path.getsize(self.db_path) / (1024 * 1024), 2)
            else:
                info['db_size_mb'] = 0

            return info

        finally:
            conn.close()

    def backup_database(self, backup_name: str = None) -> str:
        """备份数据库"""
        if backup_name is None:
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            backup_name = f"stock_data_backup_{timestamp}.db"

        backup_path = os.path.join(self.config.get_backup_path(), backup_name)

        try:
            # 使用SQLite备份API
            source_conn = self._get_connection()
            backup_conn = sqlite3.connect(backup_path)

            source_conn.backup(backup_conn)

            backup_conn.close()
            source_conn.close()

            self.logger.info(f"数据库备份成功: {backup_path}")
            return backup_path

        except Exception as e:
            self.logger.error(f"数据库备份失败: {e}")
            return None

    def optimize_database(self):
        """优化数据库性能"""
        conn = self._get_connection()
        try:
            cursor = conn.cursor()

            # 重新构建索引
            cursor.execute("REINDEX")

            # 清理空间
            cursor.execute("VACUUM")

            # 更新统计信息
            cursor.execute("ANALYZE")

            conn.commit()
            self.logger.info("数据库优化完成")

        except Exception as e:
            self.logger.error(f"数据库优化失败: {e}")
        finally:
            conn.close()

    def cleanup_old_data(self, before_date: str = None):
        """清理旧数据"""
        if before_date is None:
            # 默认保留最近2年的数据
            before_date = (datetime.now() - timedelta(days=365*2)).strftime('%Y-%m-%d')

        conn = self._get_connection()
        try:
            cursor = conn.cursor()

            # 清理日线数据
            cursor.execute(
                "DELETE FROM daily_quotes WHERE trade_date < ?",
                (before_date,)
            )
            deleted_rows = cursor.rowcount

            conn.commit()
            self.logger.info(f"清理完成，删除了 {deleted_rows} 条旧数据")

        except Exception as e:
            self.logger.error(f"数据清理失败: {e}")
        finally:
            conn.close()

    def save_stock_basic(self, data: pd.DataFrame, data_source: str):
        """保存股票基本信息"""
        if data.empty:
            return

        data['data_source'] = data_source
        data['updated_time'] = datetime.now()

        conn = self._get_connection()
        try:
            data.to_sql('stock_basic', conn, if_exists='replace', index=False)
            conn.commit()
            self.logger.info(f"保存 {len(data)} 条股票基本信息，数据源: {data_source}")
        except Exception as e:
            self.logger.error(f"保存股票基本信息失败: {e}")
        finally:
            conn.close()

    def save_daily_quotes(self, data: pd.DataFrame, data_source: str):
        """保存日线行情数据"""
        if data.empty:
            return

        data['data_source'] = data_source
        data['created_time'] = datetime.now()

        conn = self._get_connection()
        try:
            # 使用自定义方法处理重复数据
            data.to_sql('daily_quotes', conn, if_exists='append', index=False,
                        method=self._insert_or_replace)
            conn.commit()
            self.logger.info(f"保存 {len(data)} 条日线数据，数据源: {data_source}")
        except Exception as e:
            self.logger.error(f"保存日线数据失败: {e}")
        finally:
            conn.close()

    def _insert_or_replace(self, table, conn, keys, data_iter):
        """处理重复数据的插入（使用REPLACE）"""
        sql = f"REPLACE INTO {table.name} ({', '.join(keys)}) VALUES ({', '.join(['?']*len(keys))})"
        cursor = conn.cursor()
        cursor.executemany(sql, data_iter)

    def get_daily_data(self, symbol: str, start_date: str = None, end_date: str = None) -> pd.DataFrame:
        """获取指定股票的日线数据"""
        conn = self._get_connection()

        query = """
            SELECT symbol, trade_date, open_price, high_price, low_price, close_price, 
                   volume, turnover, change_rate, data_source
            FROM daily_quotes 
            WHERE symbol = ?
        """
        params = [symbol]

        if start_date:
            query += " AND trade_date >= ?"
            params.append(start_date)
        if end_date:
            query += " AND trade_date <= ?"
            params.append(end_date)

        query += " ORDER BY trade_date"

        try:
            df = pd.read_sql(query, conn, params=params)
            return df
        except Exception as e:
            self.logger.error(f"查询日线数据失败: {e}")
            return pd.DataFrame()
        finally:
            conn.close()

    def get_last_trade_date(self, symbol: str = None) -> str:
        """获取最后交易日期"""
        conn = self._get_connection()

        if symbol:
            query = "SELECT MAX(trade_date) FROM daily_quotes WHERE symbol = ?"
            params = [symbol]
        else:
            query = "SELECT MAX(trade_date) FROM daily_quotes"
            params = []

        try:
            cursor = conn.cursor()
            cursor.execute(query, params)
            result = cursor.fetchone()
            return result[0] if result[0] else "2000-01-01"
        except Exception as e:
            self.logger.error(f"获取最后交易日期失败: {e}")
            return "2000-01-01"
        finally:
            conn.close()

    def save_baostock_indicators(self, data: pd.DataFrame):
        """保存Baostock特有指标"""
        if data.empty:
            return

        conn = self._get_connection()
        try:
            data.to_sql('baostock_indicators', conn, if_exists='append',
                        index=False, method=self._insert_or_replace)
            conn.commit()
            self.logger.info(f"保存 {len(data)} 条Baostock指标数据")
        except Exception as e:
            self.logger.error(f"保存Baostock指标失败: {e}")
        finally:
            conn.close()

    def save_dividend_data(self, data: pd.DataFrame, data_source: str):
        """保存分红数据"""
        if data.empty:
            return

        data['data_source'] = data_source
        conn = self._get_connection()
        try:
            data.to_sql('dividend_data_enhanced', conn, if_exists='append',
                        index=False, method=self._insert_or_replace)
            conn.commit()
            self.logger.info(f"保存 {len(data)} 条分红数据")
        except Exception as e:
            self.logger.error(f"保存分红数据失败: {e}")
        finally:
            conn.close()